require "socket"

module Rack

  # A middleware for injecting upstream host info into a response header.
  class UpstreamIdentification

    attr_reader :app, :transform, :header_name

    # Creates the middleware with the following configuration options:
    #
    # [+header+] The name of the header that is sent to the client.
    #             Defaults to "X-Upstream-Identification"
    #
    # [+blk+]    The optional block used to transform a header. If
    #             a block is passed, the following applies:
    #               * The machine's hostname is passed to the block.
    #                 This gives the block the opportunity to change
    #                 how the hostname is rendered (obfuscate, mangle,
    #                 whatever).
    #               * The block should return a string value that will
    #                 be set as the configured header value.
    #
    #             If no block is passed during configuration, just send
    #             the hostname as a default.
    #
    def initialize(app, header = "X-Upstream-Identification", &blk)
      @app = app
      @header_name = header
      @transform = blk if block_given?
    end

    def call(env)
      status, headers, response = app.call(env)
      inject_identity! headers

      [status, headers, response]
    end

    protected

    # Push the configured header down to the client
    def inject_identity!(headers)
      headers[header_name] = if transform
        invoke_transform host_name, &transform
      else
        host_name
      end

      nil
    end

    # Invoke the transform block.
    # If the block does not return a non-nil value, raise and exception
    def invoke_transform(host_name, &transform)
      header_value = transform.call(host_name)
      unless header_value
        raise ArgumentError, "transform block must return a non-nil value"
      end

      header_value.to_s
    end

    # Get the hostname for the current machine
    def host_name
      @host_name ||= Socket.gethostname
    end

  end

end
